package com.mini.gateway.filter;

import com.mini.common.RedisService;
import com.mini.common.constant.RedisConstant;
import com.mini.common.util.JacksonUtil;
import com.mini.common.vo.ResData;
import com.mini.gateway.ApplicationContextProvider;
import org.apache.commons.lang.math.RandomUtils;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.discovery.DiscoveryClient;
import org.springframework.cloud.client.loadbalancer.LoadBalancerUriTools;
import org.springframework.cloud.client.loadbalancer.reactive.DefaultRequest;
import org.springframework.cloud.client.loadbalancer.reactive.DefaultResponse;
import org.springframework.cloud.client.loadbalancer.reactive.Request;
import org.springframework.cloud.client.loadbalancer.reactive.Response;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.ReactiveLoadBalancerClientFilter;
import org.springframework.cloud.gateway.support.DelegatingServiceInstance;
import org.springframework.cloud.gateway.support.NotFoundException;
import org.springframework.core.Ordered;
import org.springframework.http.HttpHeaders;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.netty.ByteBufFlux;

import java.net.URI;
import java.util.List;

import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.*;

/**
 * @author 小鱼儿
 * @date 2020/11/17 16:13
 */
public class ShortStateLoadBalancer implements GatewayFilter, Ordered {

    private DiscoveryClient discoveryClient;

    private RedisService redisService;

    public ShortStateLoadBalancer() {
        discoveryClient = ApplicationContextProvider.getBean(DiscoveryClient.class);
        redisService = ApplicationContextProvider.getBean(RedisService.class);
    }

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        URI url = exchange.getAttribute(GATEWAY_REQUEST_URL_ATTR);
        String schemePrefix = exchange.getAttribute(GATEWAY_SCHEME_PREFIX_ATTR);
        String lbStr = "sslb";
        if (url == null
                || (!lbStr.equals(url.getScheme()) && !lbStr.equals(schemePrefix))) {
            return chain.filter(exchange);
        }
        // preserve the original url
        addOriginalRequestUrl(exchange, url);
        URI uri = exchange.getAttribute(GATEWAY_REQUEST_URL_ATTR);
        List<ServiceInstance> list = discoveryClient.getInstances(uri.getHost());

        String token = exchange.getRequest().getHeaders().getFirst("token");
        ServerHttpResponse serverHttpResponse = exchange.getResponse();
        if (token == null) {
            ResData resData = new ResData();
            resData.setCode(1000);
            String data = JacksonUtil.toJsonString(resData);
            return serverHttpResponse.writeAndFlushWith(Flux.just(ByteBufFlux.just(serverHttpResponse.bufferFactory().wrap(data.getBytes()))));
        }
        return choose(exchange, list).doOnNext(response -> {

            if (!response.hasServer()) {
                throw NotFoundException.create(false,
                        "Unable to find instance for " + url.getHost());
            }

            URI uri2 = exchange.getRequest().getURI();

            // if the `lb:<scheme>` mechanism was used, use `<scheme>` as the default,
            // if the loadbalancer doesn't provide one.
            String overrideScheme = null;
            if (schemePrefix != null) {
                overrideScheme = url.getScheme();
            }

            DelegatingServiceInstance serviceInstance = new DelegatingServiceInstance(
                    response.getServer(), overrideScheme);

            URI requestUrl = reconstructURI(serviceInstance, uri2);

            exchange.getAttributes().put(GATEWAY_REQUEST_URL_ATTR, requestUrl);
        }).then(chain.filter(exchange));
    }

    protected URI reconstructURI(ServiceInstance serviceInstance, URI original) {
        return LoadBalancerUriTools.reconstructURI(serviceInstance, original);
    }


    private Mono<Response<ServiceInstance>> choose(ServerWebExchange exchange, List<ServiceInstance> instances) {
        return Mono.just(getInstanceResponse(instances, exchange.getRequest().getHeaders()));
    }

    public Response<ServiceInstance> getInstanceResponse(List<ServiceInstance> instances, HttpHeaders headers) {
        String token = headers.getFirst("token");
        String serverId = redisService.get(RedisConstant.TOKEN_SERVER_ID + token);
        ServiceInstance serviceInstance = null;
        if (serverId != null) {
            for (ServiceInstance e : instances) {
                String id = e.getMetadata().get("serverId");
                if (id.equals(serverId)) {
                    serviceInstance = e;
                    break;
                }
            }
        }
        if (serviceInstance == null) {
            //如果没有这个服务了 那就切换下 同时修改下token所在服务绑定
            serviceInstance = instances.get(RandomUtils.nextInt(instances.size()));
            String tempserverId = serviceInstance.getMetadata().get("serverId");
            redisService.set(RedisConstant.TOKEN_SERVER_ID + token, tempserverId, RedisConstant.TOKEN_EXPIRE_TIME);
        }
        return new DefaultResponse(serviceInstance);
    }

    @Override
    public int getOrder() {
        return 10150;
    }
}
